#include "..\CUDA_ComputeStructureTensor.h"
#include "CUDA_ST_MethodPosition.h"

__global__ void 
FiberParameterization_kernel(CUDA_Fiber *fibers, Fiber_Param *parameters, const int stepSize, const int num)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;

	int index = y * gridDim.x * blockDim.x + x;

	if (index < num) {
		/* copy the seed */
		parameters[index].seed.pos = fibers[index].seed.pos;
		parameters[index].seed.dir = fibers[index].seed.dir;

		/* parameterization */
		
		// seed
		parameters[index].parameters[PARAM_COUNT/2] = fibers[index].seed.pos;
		
		// forward
		for (int k = 1; k <= PARAM_COUNT/2; ++k) {
			int delta = k * stepSize;
			float f = (float)delta / (float)RESAMPLE;
			int lower = (int)floor(f);
			int upper = (int)ceil(f);
			float df = f - (float)lower;

			lower = lower < fibers[index].fc ? lower : (fibers[index].fc - 1);
			upper = upper < fibers[index].fc ? upper : (fibers[index].fc - 1);

			float3 v;
			v.x = fibers[index].f[upper].x - fibers[index].f[lower].x;
			v.y = fibers[index].f[upper].y - fibers[index].f[lower].y;
			v.z = fibers[index].f[upper].z - fibers[index].f[lower].z;

			parameters[index].parameters[PARAM_COUNT/2+k].x = fibers[index].f[lower].x + v.x * df;
			parameters[index].parameters[PARAM_COUNT/2+k].y = fibers[index].f[lower].y + v.y * df;
			parameters[index].parameters[PARAM_COUNT/2+k].z = fibers[index].f[lower].z + v.z * df;
		}

		// backward
		for (int k = 1; k <= PARAM_COUNT/2; ++k) {
			int delta = k * stepSize;
			float f = (float)delta / (float)RESAMPLE;
			int lower = (int)floor(f);
			int upper = (int)ceil(f);
			float df = f - (float)lower;

			lower = lower < fibers[index].bc ? lower : (fibers[index].bc - 1);
			upper = upper < fibers[index].bc ? upper : (fibers[index].bc - 1);

			float3 v;
			v.x = fibers[index].b[upper].x - fibers[index].b[lower].x;
			v.y = fibers[index].b[upper].y - fibers[index].b[lower].y;
			v.z = fibers[index].b[upper].z - fibers[index].b[lower].z;

			parameters[index].parameters[PARAM_COUNT/2-k].x = fibers[index].b[lower].x + v.x * df;
			parameters[index].parameters[PARAM_COUNT/2-k].y = fibers[index].b[lower].y + v.y * df;
			parameters[index].parameters[PARAM_COUNT/2-k].z = fibers[index].b[lower].z + v.z * df;
		}
	}
}


extern "C"
void CUDA_FiberParameterization_DTI(CFiber **fibers,
		Fiber_Param *parameters, const int w, const int h, const int stepSize)
{
	int remain_num, last_num, this_num, num_per_run, i, j;
	
	remain_num = w * h;
	last_num = 0;
	num_per_run = THREADS_PER_RUN;

	CUDA_Fiber *h_fibers = (CUDA_Fiber *)malloc(sizeof(CUDA_Fiber)*num_per_run);
	
	CUDA_Fiber *d_fibers = NULL;
	cutilSafeCall(cudaMalloc((void **)&d_fibers, sizeof(CUDA_Fiber)*num_per_run));
	Fiber_Param *d_parameters = NULL;
	cutilSafeCall(cudaMalloc((void **)&d_parameters, sizeof(Fiber_Param)*num_per_run));

	while (remain_num > 0) {
		this_num = num_per_run < remain_num ? num_per_run : remain_num;
		for (i = 0; i < this_num; ++i) {
			h_fibers[i].fc = fibers[i+last_num][0].m_fCount;
			h_fibers[i].bc = fibers[i+last_num][0].m_bCount;
			h_fibers[i].seed.pos.x = fibers[i+last_num][0].m_seed.pos.m_x;
			h_fibers[i].seed.pos.y = fibers[i+last_num][0].m_seed.pos.m_y;
			h_fibers[i].seed.pos.z = fibers[i+last_num][0].m_seed.pos.m_z;
			h_fibers[i].seed.dir.x = fibers[i+last_num][0].m_seed.dir.m_x;
			h_fibers[i].seed.dir.y = fibers[i+last_num][0].m_seed.dir.m_y;
			h_fibers[i].seed.dir.z = fibers[i+last_num][0].m_seed.dir.m_z;

			for (j = 0; j < h_fibers[i].fc; ++j) {
				h_fibers[i].f[j].x = fibers[i+last_num][0].m_pF[j].m_x;
				h_fibers[i].f[j].y = fibers[i+last_num][0].m_pF[j].m_y;
				h_fibers[i].f[j].z = fibers[i+last_num][0].m_pF[j].m_z;
			}

			for (j = 0; j < h_fibers[i].bc; ++j) {
				h_fibers[i].b[j].x = fibers[i+last_num][0].m_pB[j].m_x;
				h_fibers[i].b[j].y = fibers[i+last_num][0].m_pB[j].m_y;
				h_fibers[i].b[j].z = fibers[i+last_num][0].m_pB[j].m_z;
			}
		}

		cutilSafeCall(cudaMemcpy(d_fibers, h_fibers, 
			sizeof(CUDA_Fiber)*this_num, cudaMemcpyHostToDevice));

		dim3 dimBlock(64, 1);
		dim3 dimGrid(num_per_run/64, 1);

		FiberParameterization_kernel<<<dimGrid, dimBlock>>>(d_fibers, d_parameters, stepSize, this_num);
		cutilSafeCall(cudaGetLastError());

		cutilSafeCall(cudaMemcpy(&(parameters[last_num]), d_parameters, 
			sizeof(Fiber_Param)*this_num, cudaMemcpyDeviceToHost));

		remain_num -= this_num;
		last_num += this_num;

		printf("#");
		fflush(stdout);
	}

	cutilSafeCall(cudaFree(d_fibers));
	cutilSafeCall(cudaFree(d_parameters));

	free(h_fibers);
}

extern "C"
void CUDA_FiberParameterization_HARDI(CFiber **fibers, int *count, 
		Fiber_Param *parameters, const int w, const int h, const int stepSize)
{
	int remain_num, last_num, this_num, num_per_run, i, j, k;
	
	remain_num = w * h * 3;
	last_num = 0;
	num_per_run = THREADS_PER_RUN * 3;

	CUDA_Fiber *h_fibers = (CUDA_Fiber *)malloc(sizeof(CUDA_Fiber)*num_per_run);
	memset(h_fibers, 0, sizeof(CUDA_Fiber)*num_per_run);
	
	CUDA_Fiber *d_fibers = NULL;
	cutilSafeCall(cudaMalloc((void **)&d_fibers, sizeof(CUDA_Fiber)*num_per_run));
	Fiber_Param *d_parameters = NULL;
	cutilSafeCall(cudaMalloc((void **)&d_parameters, sizeof(Fiber_Param)*num_per_run));

	while (remain_num > 0) {
		this_num = num_per_run < remain_num ? num_per_run : remain_num;
		for (i = 0; i < this_num/3; ++i) {
			for (k = 0; k < count[i+last_num]; ++k) {
				h_fibers[i*3+k].fc = fibers[i+last_num][k].m_fCount;
				h_fibers[i*3+k].bc = fibers[i+last_num][k].m_bCount;
				h_fibers[i*3+k].seed.pos.x = fibers[i+last_num][k].m_seed.pos.m_x;
				h_fibers[i*3+k].seed.pos.y = fibers[i+last_num][k].m_seed.pos.m_y;
				h_fibers[i*3+k].seed.pos.z = fibers[i+last_num][k].m_seed.pos.m_z;
				h_fibers[i*3+k].seed.dir.x = fibers[i+last_num][k].m_seed.dir.m_x;
				h_fibers[i*3+k].seed.dir.y = fibers[i+last_num][k].m_seed.dir.m_y;
				h_fibers[i*3+k].seed.dir.z = fibers[i+last_num][k].m_seed.dir.m_z;

				for (j = 0; j < h_fibers[i*3+k].fc; ++j) {
					h_fibers[i*3+k].f[j].x = fibers[i+last_num][k].m_pF[j].m_x;
					h_fibers[i*3+k].f[j].y = fibers[i+last_num][k].m_pF[j].m_y;
					h_fibers[i*3+k].f[j].z = fibers[i+last_num][k].m_pF[j].m_z;
				}

				for (j = 0; j < h_fibers[i*3+k].bc; ++j) {
					h_fibers[i*3+k].b[j].x = fibers[i+last_num][k].m_pB[j].m_x;
					h_fibers[i*3+k].b[j].y = fibers[i+last_num][k].m_pB[j].m_y;
					h_fibers[i*3+k].b[j].z = fibers[i+last_num][k].m_pB[j].m_z;
				}
			}
		}

		cutilSafeCall(cudaMemcpy(d_fibers, h_fibers, 
			sizeof(CUDA_Fiber)*this_num, cudaMemcpyHostToDevice));

		dim3 dimBlock(64, 1);
		dim3 dimGrid(num_per_run/64, 1);

		FiberParameterization_kernel<<<dimGrid, dimBlock>>>(d_fibers, d_parameters, stepSize, this_num);
		cutilSafeCall(cudaGetLastError());

		cutilSafeCall(cudaMemcpy(&(parameters[last_num*3]), d_parameters, 
			sizeof(Fiber_Param)*this_num, cudaMemcpyDeviceToHost));

		remain_num -= this_num;
		last_num += this_num/3;

		printf("#");
		fflush(stdout);
	}

	cutilSafeCall(cudaFree(d_fibers));
	cutilSafeCall(cudaFree(d_parameters));

	free(h_fibers);
}